#pragma once

#include "viml/ast.h"
#include "viml/json_encoder.h"
#include <stack>

namespace viml {

// AST JSON打印器 - 使用Visitor模式输出JSON格式的AST
class ASTJsonPrinter : public ASTVisitor {
private:
    std::stack<std::string> json_stack;  // 用于构建嵌套JSON结构
    
    // 辅助方法：获取当前构建的JSON字符串
    std::string pop_json() {
        if (json_stack.empty()) return "null";
        auto result = json_stack.top();
        json_stack.pop();
        return result;
    }
    
    // 辅助方法：推入JSON字符串到栈
    void push_json(const std::string& json) {
        json_stack.push(json);
    }
    
public:
    ASTJsonPrinter() = default;
    
    // 获取JSON输出结果
    std::string get_json() {
        return json_stack.empty() ? "null" : json_stack.top();
    }
    
    // 清空状态，准备处理新的AST
    void reset() {
        while (!json_stack.empty()) {
            json_stack.pop();
        }
    }
    
    // ASTVisitor接口实现
    void visit(NumberLiteral& node) override {
        auto json = JSONEncoder::object({
            {"type", JSONEncoder::string_value("NumberLiteral")},
            {"value", JSONEncoder::number_value(node.value)}
        });
        push_json(json);
    }
    
    void visit(StringLiteral& node) override {
        auto json = JSONEncoder::object({
            {"type", JSONEncoder::string_value("StringLiteral")},
            {"value", JSONEncoder::string_value(node.value)}
        });
        push_json(json);
    }
    
    void visit(Identifier& node) override {
        auto json = JSONEncoder::object({
            {"type", JSONEncoder::string_value("Identifier")},
            {"name", JSONEncoder::string_value(node.name)}
        });
        push_json(json);
    }
    
    void visit(SpecialConstantNode& node) override {
        auto json = JSONEncoder::object({
            {"type", JSONEncoder::string_value("SpecialConstantNode")},
            {"constant_name", JSONEncoder::string_value(node.constant_name)}
        });
        push_json(json);
    }
    
    void visit(BinaryExpression& node) override {
        // 先处理子节点
        node.left->accept(*this);
        auto left_json = pop_json();
        
        node.right->accept(*this);
        auto right_json = pop_json();
        
        // 构建二元表达式JSON
        auto json = JSONEncoder::object({
            {"type", JSONEncoder::string_value("BinaryExpression")},
            {"operator", JSONEncoder::string_value(BinaryExpression::operator_to_string(node.op))},
            {"left", left_json},
            {"right", right_json}
        });
        push_json(json);
    }
    
    void visit(UnaryExpression& node) override {
        // 先处理操作数
        node.operand->accept(*this);
        auto operand_json = pop_json();
        
        // 构建一元表达式JSON
        auto json = JSONEncoder::object({
            {"type", JSONEncoder::string_value("UnaryExpression")},
            {"operator", JSONEncoder::string_value(UnaryExpression::operator_to_string(node.op))},
            {"operand", operand_json}
        });
        push_json(json);
    }
    
    void visit(CallExpression& node) override {
        // 处理被调用的表达式
        node.callee->accept(*this);
        auto callee_json = pop_json();
        
        // 处理参数列表
        std::vector<std::string> arg_jsons;
        for (auto& arg : node.arguments) {
            arg->accept(*this);
            arg_jsons.push_back(pop_json());
        }
        
        // 构建函数调用JSON
        auto json = JSONEncoder::object({
            {"type", JSONEncoder::string_value("CallExpression")},
            {"callee", callee_json},
            {"arguments", JSONEncoder::array(arg_jsons)}
        });
        push_json(json);
    }
    
    void visit(ListLiteral& node) override {
        // 处理列表元素
        std::vector<std::string> element_jsons;
        for (auto& element : node.elements) {
            element->accept(*this);
            element_jsons.push_back(pop_json());
        }
        
        // 构建列表字面量JSON
        auto json = JSONEncoder::object({
            {"type", JSONEncoder::string_value("ListLiteral")},
            {"elements", JSONEncoder::array(element_jsons)}
        });
        push_json(json);
    }
    
    void visit(DictLiteral& node) override {
        // 处理字典键值对
        std::vector<std::string> pair_jsons;
        for (auto& pair : node.pairs) {
            pair.first->accept(*this);
            auto key_json = pop_json();
            
            pair.second->accept(*this);
            auto value_json = pop_json();
            
            auto pair_json = JSONEncoder::object({
                {"key", key_json},
                {"value", value_json}
            });
            pair_jsons.push_back(pair_json);
        }
        
        // 构建字典字面量JSON
        auto json = JSONEncoder::object({
            {"type", JSONEncoder::string_value("DictLiteral")},
            {"pairs", JSONEncoder::array(pair_jsons)}
        });
        push_json(json);
    }
    
    void visit(IndexExpression& node) override {
        // 处理对象和索引
        node.object->accept(*this);
        auto object_json = pop_json();
        
        node.index->accept(*this);
        auto index_json = pop_json();
        
        // 构建索引访问JSON
        auto json = JSONEncoder::object({
            {"type", JSONEncoder::string_value("IndexExpression")},
            {"object", object_json},
            {"index", index_json}
        });
        push_json(json);
    }
    
    void visit(LetStatement& node) override {
        // 处理初始值
        node.value->accept(*this);
        auto value_json = pop_json();
        
        // 构建Let语句JSON
        auto json = JSONEncoder::object({
            {"type", JSONEncoder::string_value("LetStatement")},
            {"name", JSONEncoder::string_value(node.name)},
            {"value", value_json}
        });
        push_json(json);
    }
    
    void visit(EchoStatement& node) override {
        // 处理表达式
        node.expression->accept(*this);
        auto expression_json = pop_json();
        
        // 构建Echo语句JSON
        auto json = JSONEncoder::object({
            {"type", JSONEncoder::string_value("EchoStatement")},
            {"expression", expression_json}
        });
        push_json(json);
    }
    
    void visit(IfStatement& node) override {
        // 处理条件
        node.condition->accept(*this);
        auto condition_json = pop_json();
        
        // 处理then语句块
        std::vector<std::string> then_jsons;
        for (auto& stmt : node.then_body) {
            stmt->accept(*this);
            then_jsons.push_back(pop_json());
        }
        
        // 处理else语句块
        std::vector<std::string> else_jsons;
        for (auto& stmt : node.else_body) {
            stmt->accept(*this);
            else_jsons.push_back(pop_json());
        }
        
        // 构建If语句JSON
        auto json = JSONEncoder::object({
            {"type", JSONEncoder::string_value("IfStatement")},
            {"condition", condition_json},
            {"then_body", JSONEncoder::array(then_jsons)},
            {"else_body", JSONEncoder::array(else_jsons)}
        });
        push_json(json);
    }
    
    void visit(WhileStatement& node) override {
        // 处理条件
        node.condition->accept(*this);
        auto condition_json = pop_json();
        
        // 处理循环体
        std::vector<std::string> body_jsons;
        for (auto& stmt : node.body) {
            stmt->accept(*this);
            body_jsons.push_back(pop_json());
        }
        
        // 构建While语句JSON
        auto json = JSONEncoder::object({
            {"type", JSONEncoder::string_value("WhileStatement")},
            {"condition", condition_json},
            {"body", JSONEncoder::array(body_jsons)}
        });
        push_json(json);
    }
    
    void visit(ForStatement& node) override {
        // 处理可迭代对象
        node.iterable->accept(*this);
        auto iterable_json = pop_json();
        
        // 处理循环体
        std::vector<std::string> body_jsons;
        for (auto& stmt : node.body) {
            stmt->accept(*this);
            body_jsons.push_back(pop_json());
        }
        
        // 构建For语句JSON
        auto json = JSONEncoder::object({
            {"type", JSONEncoder::string_value("ForStatement")},
            {"variable", JSONEncoder::string_value(node.variable)},
            {"iterable", iterable_json},
            {"body", JSONEncoder::array(body_jsons)}
        });
        push_json(json);
    }
    
    void visit(FunctionStatement& node) override {
        // 处理函数体
        std::vector<std::string> body_jsons;
        for (auto& stmt : node.body) {
            stmt->accept(*this);
            body_jsons.push_back(pop_json());
        }
        
        // 构建参数列表JSON
        std::vector<std::string> param_jsons;
        for (const auto& param : node.parameters) {
            param_jsons.push_back(JSONEncoder::string_value(param));
        }
        
        // 构建函数定义JSON
        auto json = JSONEncoder::object({
            {"type", JSONEncoder::string_value("FunctionStatement")},
            {"name", JSONEncoder::string_value(node.name)},
            {"parameters", JSONEncoder::array(param_jsons)},
            {"body", JSONEncoder::array(body_jsons)}
        });
        push_json(json);
    }
    
    void visit(ReturnStatement& node) override {
        std::string value_json = "null";
        if (node.value) {
            node.value->accept(*this);
            value_json = pop_json();
        }
        
        // 构建Return语句JSON
        auto json = JSONEncoder::object({
            {"type", JSONEncoder::string_value("ReturnStatement")},
            {"value", value_json}
        });
        push_json(json);
    }
    
    void visit(CallStatement& node) override {
        // 处理函数调用表达式
        node.call_expr->accept(*this);
        auto call_expr_json = pop_json();
        
        // 构建Call语句JSON
        auto json = JSONEncoder::object({
            {"type", JSONEncoder::string_value("CallStatement")},
            {"call_expr", call_expr_json}
        });
        push_json(json);
    }
    
    void visit(Program& node) override {
        // 处理所有语句
        std::vector<std::string> statement_jsons;
        for (auto& stmt : node.statements) {
            stmt->accept(*this);
            statement_jsons.push_back(pop_json());
        }
        
        // 构建程序JSON
        auto json = JSONEncoder::object({
            {"type", JSONEncoder::string_value("Program")},
            {"statements", JSONEncoder::array(statement_jsons)}
        });
        push_json(json);
    }
};

} // namespace viml